import os
import json
import base64
import hashlib
import hmac
import threading
import tkinter as tk
from tkinter import filedialog, messagebox, scrolledtext
from datetime import datetime
from time import mktime
from urllib.parse import urlencode
import ssl  # 确保 ssl 被导入

import requests
import websocket
from wsgiref.handlers import format_date_time
from PIL import Image, ImageTk
from playsound import playsound  # 导入 playsound 用于播放 MP3

# ------------------------ 配置 ------------------------ #

APPId = "9952d66b"  
APISecret = "MTVkM2ExYTlhOTg2YTkzOWQxYWZjNWI2"  
APIKey = "8daf236addd09590607db513a87db1ce"  

# ------------------------ 辅助类 ------------------------ #

class AssembleHeaderException(Exception):
    def __init__(self, msg):
        self.message = msg

class Url:
    def __init__(self, host, path, schema):
        self.host = host
        self.path = path
        self.schema = schema

# ------------------------ 辅助函数 ------------------------ #

def sha256base64(data):
    sha256 = hashlib.sha256()
    sha256.update(data)
    digest = base64.b64encode(sha256.digest()).decode(encoding='utf-8')
    return digest

def parse_url(request_url):
    stidx = request_url.index("://")
    host = request_url[stidx + 3:]
    schema = request_url[:stidx + 3]
    edidx = host.index("/")
    if edidx <= 0:
        raise AssembleHeaderException("无效的请求 URL:" + request_url)
    path = host[edidx:]
    host = host[:edidx]
    u = Url(host, path, schema)
    return u

def assemble_ws_auth_url(request_url, method="POST", api_key="", api_secret=""):
    u = parse_url(request_url)
    host = u.host
    path = u.path
    now = datetime.now()
    date = format_date_time(mktime(now.timetuple()))
    signature_origin = f"host: {host}\ndate: {date}\n{method} {path} HTTP/1.1"
    signature_sha = hmac.new(api_secret.encode('utf-8'), signature_origin.encode('utf-8'),
                             digestmod=hashlib.sha256).digest()
    signature_sha = base64.b64encode(signature_sha).decode(encoding='utf-8')
    authorization_origin = f'api_key="{api_key}", algorithm="hmac-sha256", headers="host date request-line", signature="{signature_sha}"'
    authorization = base64.b64encode(authorization_origin.encode('utf-8')).decode(encoding='utf-8')
    values = {
        "host": host,
        "date": date,
        "authorization": authorization
    }
    return request_url + "?" + urlencode(values)

# ------------------------ OCR 功能 ------------------------ #

def perform_ocr(image_path, output_text_widget):
    try:
        with open(image_path, "rb") as f:
            imageBytes = f.read()
        
        url = 'https://api.xf-yun.com/v1/private/sf8e6aca1'
        
        body = {
            "header": {
                "app_id": APPId,
                "status": 3
            },
            "parameter": {
                "sf8e6aca1": {
                    "category": "ch_en_public_cloud",
                    "result": {
                        "encoding": "utf8",
                        "compress": "raw",
                        "format": "json"
                    }
                }
            },
            "payload": {
                "sf8e6aca1_data_1": {
                    "encoding": "jpg",
                    "image": base64.b64encode(imageBytes).decode('utf-8'),
                    "status": 3
                }
            }
        }
        
        request_url = assemble_ws_auth_url(url, "POST", APIKey, APISecret)
        
        headers = {'content-type': "application/json", 'host': 'api.xf-yun.com', 'app_id': APPId}
        response = requests.post(request_url, data=json.dumps(body), headers=headers)
        
        if response.status_code != 200:
            messagebox.showerror("OCR 错误", f"HTTP 错误: {response.status_code}")
            return None
        
        tempResult = response.json()
        if 'payload' not in tempResult or 'result' not in tempResult['payload']:
            messagebox.showerror("OCR 错误", f"无效的响应: {response.text}")
            return None
        
        # Base64 解码后的内容
        decoded_text = base64.b64decode(tempResult['payload']['result']['text']).decode()
        
        # 解析 JSON
        ocr_result = json.loads(decoded_text)
        
        # 提取 'content' 字段
        extracted_text = ""
        for page in ocr_result.get("pages", []):
            for line in page.get("lines", []):
                for word in line.get("words", []):
                    extracted_text += word.get("content", "")
        
        # 更新文本框
        output_text_widget.delete(1.0, tk.END)
        output_text_widget.insert(tk.END, extracted_text)
        return extracted_text
    except Exception as e:
        messagebox.showerror("OCR 异常", str(e))
        return None

# ------------------------ TTS 功能 ------------------------ #

STATUS_FIRST_FRAME = 0  # 第一帧的标识
STATUS_CONTINUE_FRAME = 1  # 中间帧标识
STATUS_LAST_FRAME = 2  # 最后一帧的标识

class Ws_Param(object):
    # 初始化
    def __init__(self, APPID, APIKey, APISecret, Text):
        self.APPID = APPID
        self.APIKey = APIKey
        self.APISecret = APISecret
        self.Text = Text

        # 公共参数(common)
        self.CommonArgs = {"app_id": self.APPID}
        # 业务参数(business)，更多个性化参数可在官网查看
        # 将 "aue" 从 "raw" 改为 "lame" 以获取 MP3 格式的音频
        self.BusinessArgs = {"aue": "lame", "auf": "audio/L16;rate=16000", "vcn": "xiaoyan", "tte": "utf8"}
        self.Data = {"status": 2, "text": base64.b64encode(self.Text.encode('utf-8')).decode("UTF8")}

    # 生成 URL
    def create_url(self):
        url = 'wss://tts-api.xfyun.cn/v2/tts'
        # 生成 RFC1123 格式的时间戳
        now = datetime.now()
        date = format_date_time(mktime(now.timetuple()))

        # 拼接字符串
        signature_origin = "host: " + "tts-api.xfyun.cn" + "\n"
        signature_origin += "date: " + date + "\n"
        signature_origin += "GET " + "/v2/tts " + "HTTP/1.1"
        # 进行 hmac-sha256 加密
        signature_sha = hmac.new(self.APISecret.encode('utf-8'), signature_origin.encode('utf-8'),
                                 digestmod=hashlib.sha256).digest()
        signature_sha = base64.b64encode(signature_sha).decode(encoding='utf-8')

        authorization_origin = f'api_key="{self.APIKey}", algorithm="hmac-sha256", headers="host date request-line", signature="{signature_sha}"'
        authorization = base64.b64encode(authorization_origin.encode('utf-8')).decode(encoding='utf-8')
        # 将请求的鉴权参数组合为字典
        v = {
            "authorization": authorization,
            "date": date,
            "host": "tts-api.xfyun.cn"
        }
        # 拼接鉴权参数，生成 URL
        url = url + '?' + urlencode(v)
        return url

def tts(text, output_audio_path, progress_callback=None):
    def on_message(ws, message):
        try:
            message = json.loads(message)
            code = message["code"]
            sid = message["sid"]
            if code != 0:
                errMsg = message.get("message", "")
                if progress_callback:
                    progress_callback(f"sid:{sid} 调用错误:{errMsg} 代码:{code}")
                ws.close()
            else:
                if 'data' in message:
                    audio = message["data"].get("audio", "")
                    status = message["data"].get("status", 0)
                    if audio:
                        with open(output_audio_path, 'ab') as f:
                            f.write(base64.b64decode(audio))
                    if status == 2:
                        if progress_callback:
                            progress_callback("TTS 完成。")
                        ws.close()
        except Exception as e:
            if progress_callback:
                progress_callback(f"接收消息但解析异常: {e}")

    def on_error(ws, error):
        if progress_callback:
            progress_callback(f"WebSocket 错误: {error}")

    def on_close(ws):
        if progress_callback:
            progress_callback("WebSocket 已关闭。")

    def on_open(ws):
        def run():
            d = {
                "common": wsParam.CommonArgs,
                "business": wsParam.BusinessArgs,
                "data": wsParam.Data,
            }
            d = json.dumps(d)
            if progress_callback:
                progress_callback("开始发送文本数据进行 TTS...")
            ws.send(d)
        threading.Thread(target=run).start()

    wsParam = Ws_Param(APPId, APIKey, APISecret, text)
    wsUrl = wsParam.create_url()
    ws = websocket.WebSocketApp(wsUrl, on_message=on_message, on_error=on_error, on_close=on_close)
    ws.on_open = on_open
    ws.run_forever(sslopt={"cert_reqs": ssl.CERT_NONE})

# ------------------------ GUI 应用程序 ------------------------ #

class OCR_TTS_App:
    def __init__(self, master):
        self.master = master
        master.title("OCR & TTS 应用程序")

        # 选择图片按钮
        self.select_image_button = tk.Button(master, text="选择图片", command=self.select_image)
        self.select_image_button.pack(pady=10)

        # 显示选中的图片
        self.image_label = tk.Label(master)
        self.image_label.pack()

        # 显示 OCR 结果的文本框
        self.text_display = scrolledtext.ScrolledText(master, wrap=tk.WORD, width=60, height=15)
        self.text_display.pack(pady=10)

        # OCR 和 TTS 处理按钮
        self.process_button = tk.Button(master, text="执行 OCR & TTS", command=self.process_ocr_tts)
        self.process_button.pack(pady=10)

        # 状态显示标签
        self.status_label = tk.Label(master, text="状态: 空闲", fg="blue")
        self.status_label.pack(pady=5)

        self.selected_image_path = None

    def select_image(self):
        file_path = filedialog.askopenfilename(
            title="选择一张图片",
            filetypes=[("图片文件", "*.jpg *.jpeg *.png *.bmp *.gif")]
        )
        if file_path:
            self.selected_image_path = file_path
            img = Image.open(file_path)
            img.thumbnail((400, 400))
            self.photo = ImageTk.PhotoImage(img)
            self.image_label.config(image=self.photo)
            self.text_display.delete(1.0, tk.END)
            self.status_label.config(text="状态: 图片已选择", fg="green")

    def process_ocr_tts(self):
        if not self.selected_image_path:
            messagebox.showwarning("未选择图片", "请先选择一张图片。")
            return

        # 处理期间禁用按钮
        self.select_image_button.config(state=tk.DISABLED)
        self.process_button.config(state=tk.DISABLED)
        self.status_label.config(text="状态: 执行 OCR...", fg="orange")

        def task():
            recognized_text = perform_ocr(self.selected_image_path, self.text_display)
            if recognized_text:
                self.status_label.config(text="状态: OCR 完成。开始 TTS...", fg="blue")
                audio_path = "output_audio.mp3"  # 修改为 MP3 格式
                # 如果存在旧的音频文件，先删除
                if os.path.exists(audio_path):
                    os.remove(audio_path)

                def tts_progress(msg):
                    self.status_label.config(text=f"状态: {msg}", fg="blue")

                # 执行 TTS
                tts(recognized_text, audio_path, progress_callback=tts_progress)

                # 等待 TTS 完成并播放音频
                def play_audio_after_tts():
                    if os.path.exists(audio_path):
                        self.status_label.config(text="状态: 播放音频中...", fg="blue")
                        playsound(audio_path)  # 使用 playsound 播放 MP3
                        self.status_label.config(text="状态: 完成", fg="green")

                # 创建一个新线程来播放音频，避免阻塞主线程
                threading.Thread(target=play_audio_after_tts).start()
            # 处理完成后重新启用按钮
            self.select_image_button.config(state=tk.NORMAL)
            self.process_button.config(state=tk.NORMAL)

        threading.Thread(target=task).start()

# ------------------------ 主程序执行 ------------------------ #

if __name__ == "__main__":
    root = tk.Tk()
    app = OCR_TTS_App(root)
    root.mainloop()
